package com.hutu.config;

import org.apache.kafka.clients.producer.Partitioner;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.utils.Utils;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 自定义分区策略
 */
public class CustomPartitioner implements Partitioner {

    private final AtomicInteger counter = new AtomicInteger(0);
    
    // 分区策略核心逻辑
    @Override
    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
        List<PartitionInfo> partitions = cluster.partitionsForTopic(topic);
        int numPartitions = partitions.size();

        // 场景1: key为空时的轮询策略
        if (key == null) {
            return counter.getAndIncrement() % numPartitions;
        } 
        // 场景2: 按业务键的哈希值分区（但添加自定义规则）
        else {
            // 自定义规则示例：将包含"VIP"的键强制分配到最后一个分区
            if (key.toString().contains("VIP")) {
                return numPartitions - 1; // 最后一个分区
            }
            // 默认哈希分区
            return Utils.toPositive(Utils.murmur2(keyBytes)) % numPartitions;
        }
    }

    @Override
    public void configure(Map<String, ?> configs) {
        // 可获取生产者配置（例如加载自定义参数）
    }

    @Override
    public void close() {
        // 清理资源
    }
}